import numpy as np
import torch
import torch.nn as nn
import pandas as pd

from scipy.io import wavfile
from scipy import signal
import soundfile

def compute_dB(waveform):
    """
    Args:
        x (numpy.array): Input waveform (#length).
    Returns:
        numpy.array: Output array (#length).
    """
    val = max(0.0, np.mean(np.power(waveform, 2)))
    dB = 10*np.log10(val+1e-4)
    return dB

class WavAugment(object):
    def __init__(self, noise_csv_path="data/noise.csv", rir_csv_path="data/rir.csv"):
        self.noise_paths = pd.read_csv(noise_csv_path)["utt_paths"].values
        self.noise_names = pd.read_csv(noise_csv_path)["speaker_name"].values
        self.rir_paths = pd.read_csv(rir_csv_path)["utt_paths"].values

        self.rir_audios = []
        for path in self.rir_paths:
            rir, sample_rate = soundfile.read(path)
            self.rir_audios.append(rir)

    def __call__(self, waveform):
        idx = np.random.randint(0, 6)
        if idx == 0:
            waveform = self.add_gaussian_noise(waveform)

        if idx == 1:
            waveform = self.add_real_noise(waveform)

        if idx == 2:
            waveform = self.reverberate(waveform)

        if idx == 3:
            waveform = self.reverberate(waveform)
            waveform = self.add_real_noise(waveform)

        if idx == 4:
            waveform = self.reverberate(waveform)
            waveform = self.add_real_noise(waveform)

        return waveform

    def add_gaussian_noise(self, waveform):
        """
        Args:
            x (numpy.array): Input waveform array (#length).
        Returns:
            numpy.array: Output waveform array (#length).
        """
        snr = np.random.uniform(low=5, high=25)
        clean_dB = compute_dB(waveform)
        noise = np.random.randn(len(waveform))
        noise_dB = compute_dB(noise)
        noise = np.sqrt(10 ** ((clean_dB - noise_dB - snr) / 10)) * noise
        waveform = (waveform + noise)
        return waveform.astype(np.int16)

    def change_volum(self, waveform):
        """
        Args:
            x (numpy.array): Input waveform array (#length).
        Returns:
            numpy.array: Output waveform array (#length).
        """
        volum = np.random.uniform(low=0.8, high=1.1)
        waveform = waveform * volum
        return waveform.astype(np.int16)

    def add_real_noise(self, waveform):
        """
        Args:
            x (numpy.array): Input length (#length).
        Returns:
            numpy.array: Output waveform array (#length).
        """
        clean_dB = compute_dB(waveform)

        idx = np.random.randint(0, len(self.noise_paths))
        sample_rate, noise = wavfile.read(self.noise_paths[idx])
        noise_name = self.noise_names[idx]

        snr = np.random.uniform(13, 25)
        noise_length = len(noise)
        audio_length = len(waveform)

        if audio_length >= noise_length:
            shortage = audio_length - noise_length
            noise = np.pad(noise, (0, shortage), 'wrap')
        else:
            start = np.random.randint(0, (noise_length-audio_length))
            noise = noise[start:start+audio_length]

        noise_dB = compute_dB(noise)
        noise = np.sqrt(10 ** ((clean_dB - noise_dB - snr) / 10)) * noise
        waveform = (waveform + noise)
        return waveform.astype(np.int16)

    def reverberate(self, waveform):
        """
        Args:
            x (numpy.array): Input length (#length).
        Returns:
            numpy.array: Output waveform array (#length).
        """
        idx = np.random.randint(0, len(self.noise_paths))
        rir = self.rir_audios[idx]

        rir_length = len(rir)
        audio_length = len(waveform)

        if audio_length < rir_length:
            start = np.random.randint(0, (rir_length-audio_length))
            rir = rir[start:start+audio_length]

        rir = rir/np.sqrt(np.sum(rir**2))
        waveform = signal.convolve(waveform, rir, mode='full')
        return waveform[:audio_length]


if __name__ == "__main__":
    aug = WavAugment()
    sample_rate, waveform = wavfile.read("input.wav")

    gaussian_noise_wave = aug.add_gaussian_noise(waveform)
    wavfile.write("gaussian_noise_wave.wav", 16000, gaussian_noise_wave.astye(np.int16))

    real_noise_wave = aug.add_real_noise(waveform)
    wavfile.write("real_noise_wave.wav", 16000, real_noise_wave.astype(np.int16))

    change_volum_wave = aug.change_volum(waveform)
    wavfile.write("change_volum_wave.wav", 16000, change_volum_wave.astype(np.int16))

    reverberate_wave = aug.reverberate(waveform)
    wavfile.write("reverberate_wave.wav", 16000, reverberate_wave.astype(np.int16))
